import pandas as pd
from datetime import datetime
import backtrader as bt
import matplotlib.pyplot as plt
from AvgStrage import SMAStrategy as avgStrage



plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置画图时中文显示
plt.rcParams["axes.unicode_minus"] = False  # 设置画图时的负号显示


# 1.数据加载
def get_data(code='600519', starttime='2017-01-01', endtime='2020-01-01'):
    df = pd.read_csv("SHFE/AG2212.csv")
    df.columns = ['symbol', 'datetime', 'open', 'high', 'low', 'close', 'volume', 'openinterest',
                  'turnover', 'settle', 'pre_settle', 'variety', 'local_symbol']
    df['datetime'] = pd.to_datetime(df['datetime'], format='%Y%m%d')
    min_datetime = df['datetime'].min()
    max_datetime = df['datetime'].max()
    print(min_datetime, max_datetime)
    df.index = pd.to_datetime(df.datetime)
    # df['openinterest'] = 0
    # 对df的数据列进行一个整合
    df = df[['open', 'high', 'low', 'close', 'volume', 'openinterest']]
    df.dropna(axis=0,how='any')
    return df


stock_df = get_data()

fromdate = datetime(2022, 11, 17)
todata = datetime(2022, 12, 15)
# 加载并读取数据源 dataname(数据来源) fromdate（date格式） todate
data = bt.feeds.PandasData(dataname=stock_df, fromdate=fromdate, todate=todata)


# 2.构建策略



# 3.策略设置
cerebro = bt.Cerebro()  # 构建大脑
cerebro.addstrategy(avgStrage)  # 添加策略
cerebro.adddata(data)  # 添加数据
# 设置初始金额
cash = 100000
cerebro.broker.setcash(cash)
# 设置手续费
cerebro.broker.setcommission(commission=0.0002)

cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

# 4.运行
f = fromdate.strftime('%Y-%m-%d')
t = todata.strftime('%Y-%m-%d')
print('初始金额: %.2f 回测时间: %s ~ %s' % (cash, f, t))
result = cerebro.run()
print('最终金额: %.2f 回测时间: %s ~ %s' % (cerebro.broker.getvalue(), f,
                                            t))
strats = result[0]
print('Sharpe Ratio:', strats.analyzers.sharpe.get_analysis())
print('DrawDown:', strats.analyzers.drawdown.get_analysis())

cerebro.plot(style='candlestick', barup='red', bardown='green', fmt_x_data='%Y-%m-%d', fmt_x_ticks='%Y-%m-%d')

